import numpy as np
import matplotlib.pyplot as plt
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
import onnxruntime as rt

# 设置随机种子以确保结果可复现
np.random.seed(42)
torch.manual_seed(42)

# 定义PyTorch版SVM模型 - 添加Softmax层
class SVMModel(nn.Module):
    def __init__(self, input_size, num_classes):
        super(SVMModel, self).__init__()
        self.linear = nn.Linear(input_size, num_classes)
        
    def forward(self, x):
        logits = self.linear(x)
        # 不在训练时添加softmax，因为CrossEntropyLoss已包含softmax
        if not self.training:
            return torch.softmax(logits, dim=1)  # 添加softmax转换概率
        return logits

# 生成模拟天气数据
def generate_weather_data(n_samples=1000):
    """
    生成模拟天气数据
    输入特征: 温度(℃), 湿度(%), 辐照度(W/m²)
    输出标签: 0-晴天, 1-多云, 2-雨天
    """
    # 生成特征
    temperature = np.random.uniform(5, 35, n_samples)  # 温度范围5-35℃
    humidity = np.random.uniform(20, 95, n_samples)    # 湿度范围20-95%
    radiation = np.random.uniform(50, 1000, n_samples) # 辐照度范围50-1000 W/m²
    
    # 创建特征矩阵
    X = np.column_stack((temperature, humidity, radiation))
    
    # 生成标签（基于一些简单规则）
    y = np.zeros(n_samples, dtype=int)
    
    for i in range(n_samples):
        temp, hum, rad = X[i]
        
        # 晴天: 温度较高，湿度较低，辐照度高
        if rad > 600 and hum < 60 and temp > 20:
            y[i] = 0
        # 雨天: 湿度高，辐照度低
        elif hum > 75 and rad < 400:
            y[i] = 2
        # 多云: 其他情况
        else:
            y[i] = 1
    
    # 添加一些随机噪声使数据更真实
    noise_idx = np.random.choice(n_samples, size=int(n_samples*0.1), replace=False)
    y[noise_idx] = np.random.randint(0, 3, size=len(noise_idx))
    
    return X, y

# 训练模型
def train_model(model, train_loader, criterion, optimizer, num_epochs=100):
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        for inputs, labels in train_loader:
            # 清零梯度
            optimizer.zero_grad()
            
            # 前向传播
            outputs = model(inputs)
            
            # 计算损失
            loss = criterion(outputs, labels)
            
            # 反向传播和优化
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
        
        # 每10个epoch打印一次损失
        if epoch % 10 == 9:
            print(f'Epoch {epoch+1}, Loss: {running_loss/len(train_loader):.4f}')

# 评估模型
def evaluate_model(model, test_loader):
    model.eval()
    correct = 0
    total = 0
    predictions = []
    true_labels = []
    
    with torch.no_grad():
        for inputs, labels in test_loader:
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            
            predictions.extend(predicted.cpu().numpy())
            true_labels.extend(labels.cpu().numpy())
    
    accuracy = correct / total
    return accuracy, np.array(predictions), np.array(true_labels)

# 主函数
def main():
    # 生成数据
    X, y = generate_weather_data(1000)
    print(f"数据形状: {X.shape}, 标签数量: {np.bincount(y)}")
    
    # 划分训练集和测试集
    from sklearn.model_selection import train_test_split
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    # 特征标准化
    from sklearn.preprocessing import StandardScaler
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)
    
    # 保存标准化器以便后续使用
    import pickle
    with open('weather_scaler.pkl', 'wb') as f:
        pickle.dump(scaler, f)
    print("标准化器已保存为 weather_scaler.pkl")
    
    # 转换为PyTorch张量
    X_train_tensor = torch.FloatTensor(X_train_scaled)
    y_train_tensor = torch.LongTensor(y_train)
    X_test_tensor = torch.FloatTensor(X_test_scaled)
    y_test_tensor = torch.LongTensor(y_test)
    
    # 创建数据加载器
    train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
    test_dataset = TensorDataset(X_test_tensor, y_test_tensor)
    
    batch_size = 32
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size)
    
    # 初始化模型
    input_size = 3  # 温度、湿度、辐照度
    num_classes = 3  # 晴天、多云、雨天
    model = SVMModel(input_size, num_classes)
    
    # 定义损失函数和优化器
    # 使用交叉熵损失而不是Hinge Loss，因为PyTorch中的多类问题更好处理
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    
    # 训练模型
    print("开始训练PyTorch模型...")
    train_model(model, train_loader, criterion, optimizer, num_epochs=100)
    
    # 评估模型
    accuracy, y_pred, y_true = evaluate_model(model, test_loader)
    print(f"模型准确率: {accuracy:.4f}")
    
    # 输出分类报告
    from sklearn.metrics import classification_report
    print("\n分类报告:")
    print(classification_report(y_true, y_pred, target_names=['晴天', '多云', '雨天']))
    
    # 可视化部分结果（2D投影）
    plt.figure(figsize=(10, 6))
    plt.scatter(X_test_scaled[:, 0], X_test_scaled[:, 1], c=y_test, cmap='viridis', 
                alpha=0.7, s=50, edgecolors='k')
    plt.title('天气预测-测试数据集 (温度 vs 湿度)')
    plt.xlabel('标准化温度')
    plt.ylabel('标准化湿度')
    plt.colorbar(label='天气类型 (0:晴天, 1:多云, 2:雨天)')
    plt.savefig('weather_visualization.png')
    plt.close()
    
    # 导出ONNX模型
    print("\n转换模型为ONNX格式...")
    dummy_input = torch.randn(1, 3)
    onnx_filename = 'weather_pytorch_model.onnx'
    
    # 确保模型处于评估模式，这样forward会应用softmax
    model.eval()
    
    # 导出模型为ONNX格式
    torch.onnx.export(
        model,
        dummy_input,
        onnx_filename,
        export_params=True,
        opset_version=12,
        do_constant_folding=True,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
    )
    print(f"ONNX模型已保存为 {onnx_filename}")
    
    # 测试ONNX模型
    print("\n验证ONNX模型...")
    session = rt.InferenceSession(onnx_filename)
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    
    # 使用ONNX运行时进行预测
    pred_onnx = []
    for i in range(len(X_test_scaled)):
        input_data = X_test_scaled[i:i+1].astype(np.float32)
        pred = session.run([output_name], {input_name: input_data})[0]
        pred_label = np.argmax(pred, axis=1)[0]
        pred_onnx.append(pred_label)
    
    # 比较ONNX模型和原始模型输出
    from sklearn.metrics import accuracy_score
    pred_onnx = np.array(pred_onnx)
    accuracy_onnx = accuracy_score(y_test, pred_onnx)
    print(f"ONNX模型准确率: {accuracy_onnx:.4f}")
    
    # 输出模型文件大小
    model_size = os.path.getsize(onnx_filename) / 1024.0
    print(f"ONNX模型大小: {model_size:.2f} KB")
    
    # 示例预测
    print("\n示例预测:")
    sample_indices = np.random.choice(len(X_test), 5, replace=False)
    for i in sample_indices:
        features = X_test[i]
        features_scaled = X_test_scaled[i:i+1].astype(np.float32)
        true_label = y_test[i]
        pred_label = pred_onnx[i]
        
        weather_types = ['晴天', '多云', '雨天']
        print(f"样本: 温度={features[0]:.1f}℃, 湿度={features[1]:.1f}%, 辐照度={features[2]:.1f}W/m²")
        print(f"预测: {weather_types[pred_label]}, 真实: {weather_types[true_label]}")
        
        # 获取概率值
        pred = session.run([output_name], {input_name: features_scaled})[0][0]
        softmax = np.exp(pred) / np.sum(np.exp(pred))
        print(f"预测概率: 晴天={softmax[0]:.4f}, 多云={softmax[1]:.4f}, 雨天={softmax[2]:.4f}")
        print("-" * 50)
    
    print("\n完成！生成的ONNX模型可以导入到STM32 Cube.AI中使用")

if __name__ == "__main__":
    main()
